from .data.table_api import ms_to_datetime_no_tz
from .data.table_orm import ExperimentDB, get_session_factory
from datetime import datetime
import ray
from ray.experimental.state.api import list_tasks
import threading
import time
from trade.data.add_factors import remote_add_batch_factors_dag
from trade.data.watch_kline_data import remote_watch_multi_factors, remote_watch_ohlcv


class TaskControl:
    def __init__(self, consider_task_names: dict[str, ray.ObjectRef] = None):
        self.session_factory = get_session_factory("db_exp")
        self.task_id_map_ref = {}
        self.consider_task_names = consider_task_names if consider_task_names else {}
        self.th = threading.Thread(target=self.sync_ray_info)
        self.th.start()

    def read_exp_info(
        self,
        expid_list: list[str] = None,
        title_starts: str = None,
        name_list: list[str] = None,
        status_list: list[str] = None,
        create_start: int = None,
        create_end: int = None,
    ):
        with self.session_factory() as session:
            with session.begin():
                query = session.query(
                    ExperimentDB.expid,
                    ExperimentDB.title,
                    ExperimentDB.name,
                    ExperimentDB.create,
                    ExperimentDB.start,
                    ExperimentDB.end,
                    ExperimentDB.status,
                    ExperimentDB.info,
                )
                if expid_list:
                    query = query.filter(ExperimentDB.expid.in_(expid_list))
                if title_starts:
                    query = query.filter(ExperimentDB.title.startswith(title_starts))
                if name_list:
                    query = query.filter(ExperimentDB.name.in_(name_list))
                if status_list:
                    query = query.filter(ExperimentDB.status.in_(status_list))
                if create_start is not None:
                    create_start = ms_to_datetime_no_tz(create_start)
                    query = query.filter(ExperimentDB.create >= create_start)
                if create_end is not None:
                    create_end = ms_to_datetime_no_tz(create_end)
                    query = query.filter(ExperimentDB.create <= create_end)

                results = query.all()
                results = [x._asdict() for x in results]
                fc = lambda x: (
                    int(x.timestamp() * 1000) if isinstance(x, datetime) else x
                )
                for i in range(len(results)):
                    for y in ["create", "start", "end"]:
                        results[i][y] = fc(results[i][y])
                return results

    def read_exp_result(self, expid: str):
        with self.session_factory() as session:
            with session.begin():
                query = session.query(ExperimentDB.results)
                rs = query.filter(ExperimentDB.expid == expid).first()
                return rs.results if rs else None

    def set_exp_data(
        self,
        expid,
        title=None,
        name=None,
        create=None,
        start=None,
        end=None,
        status=None,
        info=None,
        info_kv=None,
        results=None,
    ):
        updates = {}
        if title is not None:
            updates[ExperimentDB.title.name] = title
        if name is not None:
            updates[ExperimentDB.name.name] = name
        if create is not None:
            create = ms_to_datetime_no_tz(create)
            updates[ExperimentDB.create.name] = create
        if start is not None:
            start = ms_to_datetime_no_tz(start)
            updates[ExperimentDB.start.name] = start
        if end is not None:
            end = ms_to_datetime_no_tz(end)
            updates[ExperimentDB.end.name] = end
        if status is not None:
            updates[ExperimentDB.status.name] = status
        if info is not None:
            updates[ExperimentDB.info.name] = info

        if results is not None:
            updates[ExperimentDB.results.name] = results

        with self.session_factory() as session:
            with session.begin():
                record = (
                    session.query(ExperimentDB)
                    .filter(ExperimentDB.expid == expid)
                    .first()
                )
                if record:
                    if info is None and info_kv is not None:
                        updates[ExperimentDB.info] = record.info
                        for key in info_kv:
                            updates[ExperimentDB.info][key] = info_kv[key]
                    session.query(ExperimentDB).filter(
                        ExperimentDB.expid == expid
                    ).update(updates)
                else:
                    new_record = ExperimentDB(expid=expid, **updates)
                    session.add(new_record)

    def add_task(self, task_name, title, *args, **kwargs):
        if task_name in self.consider_task_names:
            obj = self.consider_task_names[task_name].remote(*args, **kwargs)
            tid = obj.task_id().hex()
            self.task_id_map_ref[tid] = obj
            self.set_exp_data(
                expid=tid,
                title=title,
                name=task_name,
                create=int(time.time() * 1000),
                status="ADD",
                info={"args": args, "kwargs": kwargs},
            )
            return tid

    def cancel_task(self, task_id):
        if task_id in self.task_id_map_ref:
            ray.cancel(self.task_id_map_ref[task_id], force=True)
            self.set_exp_data(
                expid=task_id, end=int(time.time() * 1000), status="MANUAL_CANCELLED"
            )
            del self.task_id_map_ref[task_id]
            return True
        return False

    def set_job_func(self, func, name):
        self.consider_task_names[name] = func

    def delete_job_func(self, name: str):
        del self.consider_task_names[name]

    def sync_ray_info(self):
        while True:
            time.sleep(2)
            tasks_info_list = []
            names = list(self.consider_task_names.keys())
            for x in names:
                tasks_info_list.extend(list_tasks(filters=[("name", "=", x)]))
            for tinfo in tasks_info_list:
                if tinfo.task_id not in self.task_id_map_ref:
                    continue

                print(f"fresh task state {tinfo.task_id}")
                if tinfo.state == "FINISHED":
                    rs = ray.get(self.task_id_map_ref[tinfo.task_id])
                    self.set_exp_data(
                        expid=tinfo.task_id,
                        title=None,
                        start=tinfo.start_time_ms,
                        end=tinfo.end_time_ms,
                        status=tinfo.state,
                        info_kv={
                            "error_message": tinfo.error_message,
                            "task_log_info": tinfo.task_log_info,
                        },
                        results=rs,
                    )
                    del self.task_id_map_ref[tinfo.task_id]
                else:
                    self.set_exp_data(
                        expid=tinfo.task_id,
                        title=None,
                        start=tinfo.start_time_ms,
                        end=tinfo.end_time_ms,
                        status=tinfo.state,
                        info_kv={
                            "error_message": tinfo.error_message,
                            "task_log_info": tinfo.task_log_info,
                        },
                    )


@ray.remote(
    max_restarts=3,
    max_task_retries=1,
    resources={"task_trace": 1},
)
class RemoteTaskControl(TaskControl):
    pass


def get_task_control_obj():
    name = "task_control"
    space = "tradecoin"
    all_namespaces = ray.util.list_named_actors(True)
    for x in all_namespaces:
        if x["name"] == name and x["namespace"] == space:
            obj = ray.get_actor(name=name, namespace=space)
            return obj
    obj = RemoteTaskControl.options(
        name=name, lifetime="detached", namespace=space
    ).remote()
    return obj


@ray.remote
def work_to_cost_time(x, *args, **kwargs):
    print(f"cost some time with {args} {kwargs}")
    time.sleep(x)
    print(f"done some job")
    return {"results": (args, kwargs)}


task_doc_list = [
    {
        "task_name": "work_to_cost_time",
        "func": work_to_cost_time,
        "template_params": {
            "x": 10,
            "bfname_list": ["kline"],
            "fname_list": ["open", "close"],
            "symbol_list": ["BTC/USDT"],
            "start": 1747767885189,
            "end": 1747767885189,
            "timeframe": "1m",
            "max_part": 1000,
            "n_jobs": 2,
            "is_cover": False,
        },
    },
    {
        "task_name": "remote_add_batch_factors_dag",
        "func": remote_add_batch_factors_dag,
        "template_params": {
            "bfname_list": ["kline"],
            "fname_list": ["open", "close"],
            "symbol_list": ["BTC/USDT"],
            "start": 1747767885189,
            "end": 1747767885189,
            "timeframe": "1m",
            "max_part": 1000,
            "n_jobs": 2,
            "is_cover": False,
        },
    },
    {
        "task_name": "remote_watch_ohlcv",
        "func": remote_watch_ohlcv,
        "template_params": {
            "symbolsAndTimeframes": [("BTC/USDT", "1m")],
            "is_trade": True,
        },
    },
    {
        "task_name": "remote_watch_multi_factors",
        "func": remote_watch_multi_factors,
        "template_params": {
            "consider_symbols": ["BTC/USDT", "ETH/USDT"],
            "consider_fnames": ["open", "low"],
            "timeframe": "1m",
        },
    },
]
# consider_symbols: list[str],
# consider_fnames: list[str],
# timeframe: TimeFrameType,

# TaskState(task_id='016f65e6f1721b00fd73e7801fb47b964d48900e01000000', attempt_number=0,
# name='TableEdit.get_distinct_symbols', state='FINISHED', job_id='01000000',
# actor_id='fd73e7801fb47b964d48900e01000000', type='ACTOR_TASK',
# func_or_class_name='TableEdit.get_distinct_symbols',
# parent_task_id='ffffffffffffffffffffffffffffffffffffffff02000000',
# node_id='d83e8cc34106e90ad556530cef9f1e89dce563c589a18d1ea14b955c',
# worker_id='d84e19322818c9e9aefcaa27fba49bac025908d78d3ada80db5309f5',
# worker_pid=40468, error_type=None, language=None, required_resources=None,
# runtime_env_info=None, placement_group_id=None, events=None,
# profiling_data=None, creation_time_ms=None, start_time_ms=None, end_time_ms=None,
# task_log_info=None, error_message=None, is_debugger_paused=None)
